.section .text.entry
.global _start
_start:
    # Set up stack pointer
    la sp, stack_top
    
    # Set up interrupt vector
    la t0, trap_vector
    csrw mtvec, t0
    
    # Jump to C code
    call kernel_main
    
    # If kernel_main returns, enter infinite loop
1:  j 1b

.section .text
.global trap_vector
.align 4
trap_vector:
    # Save all registers
    addi sp, sp, -256
    sd ra, 0(sp)
    sd t0, 8(sp)
    sd t1, 16(sp)
    sd t2, 24(sp)
    sd a0, 32(sp)
    sd a1, 40(sp)
    sd a2, 48(sp)
    sd a3, 56(sp)
    sd a4, 64(sp)
    sd a5, 72(sp)
    sd a6, 80(sp)
    sd a7, 88(sp)
    sd s0, 96(sp)
    sd s1, 104(sp)
    
    # Call interrupt handler
    call handle_interrupt
    
    # Restore registers
    ld ra, 0(sp)
    ld t0, 8(sp)
    ld t1, 16(sp)
    ld t2, 24(sp)
    ld a0, 32(sp)
    ld a1, 40(sp)
    ld a2, 48(sp)
    ld a3, 56(sp)
    ld a4, 64(sp)
    ld a5, 72(sp)
    ld a6, 80(sp)
    ld a7, 88(sp)
    ld s0, 96(sp)
    ld s1, 104(sp)
    addi sp, sp, 256
    
    # Return from interrupt
    mret

.section .text
.global task_switch
.align 4
task_switch:
    # a0 = &prev->sp
    # a1 = &next->sp
    
    # 保存当前上下文
    addi sp, sp, -112
    sd ra, 0(sp)
    sd s0, 8(sp)
    sd s1, 16(sp)
    sd s2, 24(sp)
    sd s3, 32(sp)
    sd s4, 40(sp)
    sd s5, 48(sp)
    sd s6, 56(sp)
    sd s7, 64(sp)
    sd s8, 72(sp)
    sd s9, 80(sp)
    sd s10, 88(sp)
    sd s11, 96(sp)
    
    # 保存sp
    sd sp, (a0)
    
    # 加载新任务的sp
    ld sp, (a1)
    
    # 恢复上下文
    ld ra, 0(sp)
    ld s0, 8(sp)
    ld s1, 16(sp)
    ld s2, 24(sp)
    ld s3, 32(sp)
    ld s4, 40(sp)
    ld s5, 48(sp)
    ld s6, 56(sp)
    ld s7, 64(sp)
    ld s8, 72(sp)
    ld s9, 80(sp)
    ld s10, 88(sp)
    ld s11, 96(sp)
    addi sp, sp, 112
    
    ret
